package elasticsearch;

import com.alibaba.fastjson.JSON;
import com.google.common.collect.Sets;
import org.apache.http.HttpHost;
import org.elasticsearch.action.admin.indices.get.GetIndexRequest;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.search.*;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.*;
import org.elasticsearch.search.Scroll;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.time.LocalDateTime;
import java.time.ZoneOffset;
import java.util.*;
import java.util.stream.Stream;

/**
 * ES客户端，增删查改:https://www.cnblogs.com/ginb/p/8716485.html<br/>
 * Created on : 2019-07-01 20:17
 * @author lizebin
 */
public class ElasticSearchClient {

    private static Logger logger = LoggerFactory.getLogger(ElasticSearchClient.class);

    public static final String INDEX_NAME = "test_index";

    public static final String TYPE_NAME = "test_type";

    public static final RestHighLevelClient CLIENT = new RestHighLevelClient(
        RestClient.builder(
            new HttpHost("192.168.56.104", 9200, "http")));

    public static void main(String[] args) throws IOException {

        System.out.println(JSON.toJSONString(get("1562792104")));

        //insert
        /*
        Map<String, Object> data = new HashMap<>();
        data.put("date", new Date());
        data.put("long", new Long(1));
        data.put("integer", new Integer(1));
        data.put("string", "abcd");
        data.put("short", new Short("1"));
        data.put("byte", new Byte("1"));
        insert(data);
        */

        //update
        //Map<String, Object> data = new HashMap<>();
        //data.put("integer", 2);
        //update("11", data);

        //delete
        //delete("11");

        /*
        List<Map<String, Object>> list = new LinkedList<>();
        for(int i=0; i<5; i++) {
            Map<String, Object> map = new HashMap<>();
            map.put("i", i);
            list.add(map);
        }
        bulkUpdate(list);
        */

        //System.out.println(JSON.toJSONString(searchById(Sets.newHashSet("4", "2"))));

        System.out.println(JSON.toJSONString(scrollQuery()));
        CLIENT.close();
    }

    /**
     * 判断索引是否存在
     * @param index
     * @return
     * @throws IOException
     */
    public static boolean isIndexExists(String index) throws IOException {
        GetIndexRequest request = new GetIndexRequest();
        request.indices(index);
        request.local(false);
        request.humanReadable(true);
        return CLIENT.indices().exists(request);
    }

    /**
     * 根据id获取
     * @param id
     * @return
     */
    public static Map<String, Object> get(String id) throws IOException {
        /*
        * {
            "exists": true,
            "fields": { },
            "fragment": false,
            "id": "12",
            "index": "test_index",
            "source": {
                "test_field": "test1"
            },
            "sourceAsBytes": "eyJ0ZXN0X2ZpZWxkIjoidGVzdDEifQ==",
            "sourceAsBytesRef": {
                "childResources": [ ],
                "fragment": true
            },
            "sourceAsMap": {
                "$ref": "$.source"
            },
            "sourceAsString": "{\"test_field\":\"test1\"}",
            "sourceEmpty": false,
            "sourceInternal": {
                "$ref": "$.sourceAsBytesRef"
            },
            "type": "test_type",
            "version": 1
        }
        *
        * */
        GetRequest request = new GetRequest(
            INDEX_NAME,
            TYPE_NAME,
            id);
        request.realtime(true);
        GetResponse response = CLIENT.get(request);
        return response.getSource();
    }

    /**
     * 插入数据
     * @param data
     */
    public static void insert(Map<String, Object> data) throws IOException {
        String id = Objects.toString(LocalDateTime.now().toEpochSecond(ZoneOffset.UTC));
        System.out.println("insert id = " + id);
        UpdateRequest request = new UpdateRequest(INDEX_NAME, TYPE_NAME, id);
        request.doc(JSON.toJSONString(data), XContentType.JSON);
        request.scriptedUpsert(true);
        request.docAsUpsert(true);
        CLIENT.update(request);
    }

    /**
     * 插入数据
     * @param id
     * @param data
     */
    public static void update(String id, Map<String, Object> data) throws IOException {
        UpdateRequest request = new UpdateRequest(INDEX_NAME, TYPE_NAME, id);
        request.id(id);
        request.doc(JSON.toJSONString(data), XContentType.JSON);
        request.scriptedUpsert(true);
        request.docAsUpsert(true);
        CLIENT.update(request);
    }

    /**
     * 删除文档
     * @param id
     */
    public static void delete(String id) throws IOException {
        DeleteRequest request = new DeleteRequest(INDEX_NAME, TYPE_NAME, id);
        DeleteResponse response = CLIENT.delete(request);
    }

    /**
     * 批量更新，不存在则插入
     * @param list
     */
    public static void bulkUpdate(List<Map<String, Object>> list) throws IOException {
        BulkRequest request = new BulkRequest();
        list.forEach(data -> {
            request.add(new IndexRequest(INDEX_NAME, TYPE_NAME, Objects.toString(Math.abs(System.nanoTime()))).source(data, XContentType.JSON));
        });
        CLIENT.bulk(request);
    }

    /**
     * 查询所有记录
     * @return
     */
    public static List<Map<String, Object>> searchAll() throws IOException {
        SearchRequest searchRequest = new SearchRequest(INDEX_NAME);
        searchRequest.types(TYPE_NAME);
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.query(QueryBuilders.matchAllQuery());
        searchRequest.source(searchSourceBuilder);
        SearchResponse response = CLIENT.search(searchRequest);
        SearchHit[] getHits = response.getHits().getHits();
        if (getHits.length > 0) {
            List<Map<String, Object>> list = new LinkedList<>();
            Stream.of(getHits).forEach(hit -> {
                Map<String, Object> map = hit.getSourceAsMap();
                map.put("id", hit.getId());
                list.add(map);
            });
            return list;
        }
        return Collections.emptyList();
    }

    /**
     * 根据id搜索
     * @param ids
     * @return
     * @throws IOException
     */
    public static List<Map<String, Object>> searchById(Set<String> ids) throws IOException {
        if (Objects.isNull(ids) || ids.isEmpty()) {
            return Collections.emptyList();
        }
        BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery();
        ids.forEach(id -> {
            TermQueryBuilder termQueryBuilder = new TermQueryBuilder("i", id);
            //等价于:id in ("", "", "")
            boolQueryBuilder.should(termQueryBuilder);
        });
        SearchRequest searchRequest = new SearchRequest(INDEX_NAME);
        searchRequest.types(TYPE_NAME);
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.query(boolQueryBuilder);
        searchRequest.source(searchSourceBuilder);
        SearchResponse response = CLIENT.search(searchRequest);
        SearchHit[] getHits = response.getHits().getHits();
        if (getHits.length > 0) {
            List<Map<String, Object>> list = new LinkedList<>();
            Stream.of(getHits).forEach(hit -> {
                Map<String, Object> map = hit.getSourceAsMap();
                map.put("id", hit.getId());
                list.add(map);
            });
            return list;
        }
        return Collections.emptyList();
    }

    /**
     * 滚动查询
     * @return
     */
    public static List<Map<String, Object>> scrollQuery() throws IOException {

        /*
        这是控制ES滚动查询的记录容量大小，如果设置成比实际数据容量小，会报错
        put /test_index/_settings
        {
          "index":{"max_result_window":1}
        }
        */

        final Scroll scroll = new Scroll(TimeValue.timeValueMinutes(5L));
        SearchRequest searchRequest = new SearchRequest(INDEX_NAME);
        searchRequest.scroll(scroll);
        searchRequest.types(TYPE_NAME);
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery();
        TermQueryBuilder termQueryBuilder = new TermQueryBuilder("i", "4");
        boolQueryBuilder.should(termQueryBuilder);
        searchSourceBuilder.query(boolQueryBuilder);
        searchRequest.source(searchSourceBuilder);
        SearchResponse searchResponse = CLIENT.search(searchRequest);
        String scrollId = searchResponse.getScrollId();
        System.out.println("scrollId = " + scrollId);
        SearchHit[] searchHits = searchResponse.getHits().getHits();

        List<Map<String, Object>> list = new LinkedList<>();
        while (searchHits != null && searchHits.length > 0) {
            Stream.of(searchHits).forEach(hit -> {
                Map<String, Object> map = hit.getSourceAsMap();
                map.put("id", hit.getId());
                list.add(map);
            });
            SearchScrollRequest scrollRequest = new SearchScrollRequest(scrollId);
            scrollRequest.scroll(scroll);
            searchResponse = CLIENT.searchScroll(scrollRequest);
            scrollId = searchResponse.getScrollId();
            System.out.println("scrollId = " + scrollId);
            searchHits = searchResponse.getHits().getHits();
        }
        return list;
    }
}
